# coding : UTF-8
"""
作者：BingBO   时间：2022年09月10日
自动调整代码格式 ：Alt+Ctrl+L
"""
import cv2 as cv
import numpy as np
import math
from matplotlib import pyplot as plt


# 灰度重心法提取激光条纹重心函数
def gray_center_method(img):
    print("input img.shape: ", img.shape)
    height, width, channels = img.shape
    b, g, r = cv.split(img)

    ret1, th1 = cv.threshold(r, 200, 255, cv.THRESH_BINARY)  # ret1分割阈值
    th1 = cv.medianBlur(th1, 3)  # 中值滤波
    cv.imshow('threshold img', th1)

    CP_row, CP_col, Zero_col = [], [], []
    for j in range(width):  # 遍历每一列
        x0, y0 = 0, 0
        for i in range(height):  # 遍历每一行
            x0 += th1[i, j]
            y0 += th1[i, j] * (i + 1)
        if x0 != 0:
            x0_float, y0_float = float(x0), float(y0)  # 整数相除，小数直接抛弃
            point_row = round(y0_float / x0_float) - 1  #
            # 列表存储条纹中心坐标
            CP_row += [point_row]
            CP_col += [j]  # 中心点列索引列表(包含空白区域列坐标的突变)
        else:
            Zero_col += [j]  # 缺陷点的列索引(包括了交叉光条右侧 的缺少光条空白区域)

    print("CP_row:", CP_row)
    print("CP_col:", CP_col)
    print("Zero_col:", Zero_col)

    # 创建新图像，显示中心点
    # img_new = np.zeros([height, width, 3], img.dtype)  # 先创建一个单通道图片
    img_new = np.zeros_like(img)
    for i in range(0, len(CP_row)):
        img[CP_row[i], CP_col[i]] = [255, 0, 0]  # 蓝色点显示
        img_new[CP_row[i], CP_col[i]] = [255, 255, 255]
    cv.imshow("CentraPoint_Binary:", img_new)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\centralPoint.jpg", img_new)

    gapNums = []
    count = 0
    for i in range(0, len(CP_col) - 1):  # 中心点的列总数len(CP_col)，索引是0~len(CP_col)-1，下语句后面减去前面，需再减1
        if (CP_col[i + 1] - CP_col[i]) > 20:  # 中心点列坐标的突变超过阈值50，认为是边缘大断点处，排除偶尔光条的间断
            # 在图像上标记左右两断点
            img[CP_row[i], CP_col[i]] = [0, 255, 0]
            img[CP_row[i + 1], CP_col[i + 1]] = [0, 255, 0]
            # (左断点(含中心点)列坐标 , 右断点(含中心点)列坐标 , 左右断点间距)
            gapNums.append((CP_col[i], CP_col[i + 1], (CP_col[i + 1] - CP_col[i] - 1)))
            index = i  # 断裂处的列索引
            count += 1

    cv.imshow('BreakPoint_V', img)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\BreakPoint_V.png", img)  # 保存图像
    print("共发现竖直方向的间断数：", count)
    print("gapNums Info:", gapNums)

    # 取断点左边点集
    CP_row_Left = CP_row[:index + 1]
    CP_col_Left = CP_col[:index + 1]
    # print("CP_row_Left:", CP_row_Left)
    # print("CP_col_Left:", CP_col_Left)
    for i, gapNum in enumerate(gapNums, start=1):
        print("列索引%d(左断点列坐标 , 右断点列坐标，断点间距): " % i, gapNum)

    return img_new, CP_row_Left, CP_col_Left


# 对中心点图像，霍夫线变换
def HoughLines(img, CP_row_Left, CP_col_Left):
    print("**********************HoughLines**********************")
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    lines = cv.HoughLines(binary, 1, np.pi / 180, 150)  # 返回一个：math:(rho，theta)值的数组
    print("检测到直线数量：", lines.shape[0])  # (1,1,2)
    print("lines:", lines)

    for line in lines:
        rho, theta = line[0]  # rho 是距图像左上角 (0,0) 的距离，theta 为弧度表示的线旋转角度。
        print("主线角度：", theta * 180 / 3.14)

        cos, sin = np.cos(theta), np.sin(theta)
        b = rho / sin
        k = np.tan(theta - 3.1416 / 2)
        print("k , b:", k, b)

        x1, x2 = CP_col_Left[0], CP_col_Left[len(CP_col_Left) - 1]
        y1, y2 = int(b + k * x1), int(b + k * x2)

        cv.line(img, (x1, y1), (x2, y2), (0, 255, 0), 1)
        cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\HoughLines.png", img)  # 保存图像

    dMax, index = 0, 0
    for i in range(0, len(CP_row_Left)):
        # 主线方程：kx-y+b=0，(A=k B=-1 C=b)
        # d = round(math.fabs(k * CP_col_Left[i] - CP_row_Left[i] + b) / math.sqrt(1 + k * k))
        d = math.fabs(k * CP_col_Left[i] - CP_row_Left[i] + b) / math.sqrt(1 + k * k)
        if d >= dMax:
            dMax = d
            coor = (CP_row_Left[i], CP_col_Left[i])
            index = i
    print("距离主线最远中心点坐标：", coor)
    img[coor[0], coor[1]] = [0, 0, 255]
    # cv.circle(img, [coor[1], coor[0]], 1, (0, 0, 255), -1)
    cv.imshow("furthestPoint", img)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\furthestPoint.png", img)  # 保存图像

    # 分割最远点左右点集
    furthestPoint_row_Left = CP_row_Left[:index + 1]
    furthestPoint_col_Left = CP_col_Left[:index + 1]
    print("furthestPoint_row_Left:", furthestPoint_row_Left)
    print("furthestPoint_col_Left:", furthestPoint_col_Left)

    furthestPoint_row_Right = CP_row_Left[index + 1:]
    furthestPoint_col_Right = CP_col_Left[index + 1:]
    print("furthestPoint_row_Right:", furthestPoint_row_Right)
    print("furthestPoint_col_Right:", furthestPoint_col_Right)

    # 创建单通道左中心点图像
    img_new_Left = np.zeros_like(img)
    for i in range(0, len(furthestPoint_row_Left)):
        img_new_Left[furthestPoint_row_Left[i], furthestPoint_col_Left[i]] = [255, 255, 255]
    cv.imshow("img_new_Left", img_new_Left)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\img_new_Left.jpg", img_new_Left)

    # 创建单通道右中心点图像
    img_new_Right = np.zeros_like(img)
    for i in range(0, len(furthestPoint_row_Right)):
        img_new_Right[furthestPoint_row_Right[i], furthestPoint_col_Right[i]] = [255, 255, 255]
    cv.imshow("img_new_Right:", img_new_Right)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\img_new_Left.jpg", img_new_Left)

    return img_new_Left, img_new_Right, furthestPoint_row_Left, furthestPoint_col_Left, furthestPoint_row_Right, furthestPoint_col_Right


def HoughLines_Left(img_new_Left, furthestPoint_row_Left, furthestPoint_col_Left):
    # ***************1. 对左半段单通道图像hough直线拟合******************
    print("\n********************HoughLines_Left************************")
    gray = cv.cvtColor(img_new_Left, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    linesLeft = cv.HoughLines(binary, 1, np.pi / 180, 100)  # 返回一个：math:(rho，theta)值的数组
    print(linesLeft.shape)  # (1,1,2)
    print("linesLeft(rho，theta):", linesLeft)

    for lineLeft in linesLeft:
        rho, theta = lineLeft[0]
        print("左半部主线角度：", theta * 180 / 3.14)

        cos, sin = np.cos(theta), np.sin(theta)
        b = rho / sin
        k = np.tan(theta - 3.1416 / 2)
        print("左半部主线k , b:", k, b)

        x1, x2 = furthestPoint_col_Left[0], furthestPoint_col_Left[len(furthestPoint_col_Left) - 1]
        y1, y2 = int(b + k * x1), int(b + k * x2)

        cv.line(img_new_Left, (x1, y1), (x2, y2), (0, 255, 0), 1)
    cv.imshow("mainline of img_new_Left", img_new_Left)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\img_new_Left.png", img_new_Left)

    # 分割出 距离主线>阈值T的  点集
    T, index, count = 2, 0, 0
    pointOnGrooveLeft = np.zeros_like(img_new_Left)
    for i in range(0, len(furthestPoint_row_Left)):
        d = math.fabs(k * furthestPoint_col_Left[i] - furthestPoint_row_Left[i] + b) / math.sqrt(1 + k * k)
        # print("d:", d)
        if d >= T:
            pointOnGrooveLeft[furthestPoint_row_Left[i], furthestPoint_col_Left[i]] = [255, 255, 255]
            count += 1
            # coor = (furthestPoint_row_Left[i], furthestPoint_col_Left[i])
            # index = i
    print("共找到大于阈值T的点数：", count)
    cv.imshow("pointOnGrooveLeft", pointOnGrooveLeft)
    print("pointOnGrooveLeft：", pointOnGrooveLeft.shape)

    # ***************2.对左侧坡口上的图片点再进行Hough直线检测，返回该直线参数******************
    gray = cv.cvtColor(pointOnGrooveLeft, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    linesOnGrooveLeft = cv.HoughLines(binary, 1, np.pi / 180, 30)  # 返回一个：math:(rho，theta)值的数组
    print(linesOnGrooveLeft.shape)  # (1,1,2)
    print("linesOnGrooveLeft(rho，theta):", linesOnGrooveLeft)  # 获取主线参数1

    for lineLeft in linesOnGrooveLeft:
        rho, theta = lineLeft[0]
        print("左半部坡口主线角度：", theta * 180 / 3.14)

        a, b = np.cos(theta), np.sin(theta)
        x0, y0 = a * rho, b * rho  # 原点到直线的 垂点坐标
        lenL, lenR = 500, 600
        x1, y1 = int(x0 + lenL * (-b)), int(y0 + lenL * (a))
        x2, y2 = int(x0 - lenR * (-b)), int(y0 - lenR * (a))
        # 获取主线参数2
        k = (y2 - y1) / (x2 - x1)
        b = y1 - k * x1
        print("左半部坡口斜率k: %f" % k, "截距b：%f" % b)
        cv.line(img_new_Left, (x1, y1), (x2, y2), (0, 255, 0), 1)
    # cv.imshow("mainline of pointOnGrooveLeft", pointOnGrooveLeft)  # 显示不了主线，则取到主线参数也可以
    return linesOnGrooveLeft, k, b


def HoughLines_Right(img_new_Right, furthestPoint_row_Right, furthestPoint_col_Right):
    # ***************1.对右半段单通道图像hough直线拟合******************
    print("\n*********************HoughLines_Right***********************")
    gray = cv.cvtColor(img_new_Right, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    linesRight = cv.HoughLines(binary, 1, np.pi / 180, 100)  # 返回一个：math:(rho，theta)值的数组
    print(linesRight.shape)  # (1,1,2)
    print("linesRight(rho，theta):", linesRight)

    for lineRight in linesRight:
        rho, theta = lineRight[0]
        print("右半部主线角度：", theta * 180 / 3.14)

        cos, sin = np.cos(theta), np.sin(theta)
        b = rho / sin
        k = np.tan(theta - 3.1416 / 2)
        print("右半部主线k , b:", k, b)

        x1, x2 = furthestPoint_col_Right[0], furthestPoint_col_Right[len(furthestPoint_col_Right) - 1]
        y1, y2 = int(b + k * x1), int(b + k * x2)

        cv.line(img_new_Right, (x1, y1), (x2, y2), (0, 255, 0), 1)
    cv.imshow("mainline of img_new_Right", img_new_Right)
    # cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\img_new_Right.png", img_new_Right)

    # 分割出 距离主线>阈值T的  点集
    T, index, count = 2, 0, 0
    pointOnGrooveRight = np.zeros_like(img_new_Right)
    for i in range(0, len(furthestPoint_row_Right)):
        d = math.fabs(k * furthestPoint_col_Right[i] - furthestPoint_row_Right[i] + b) / math.sqrt(1 + k * k)
        # print("d:", d)
        if d >= T:
            pointOnGrooveRight[furthestPoint_row_Right[i], furthestPoint_col_Right[i]] = [255, 255, 255]
            count += 1
            # coor = (furthestPoint_row_Left[i], furthestPoint_col_Left[i])
            # index = i
    print("右坡口共找到大于阈值T的点数：", count)
    cv.imshow("pointOnGrooveRight", pointOnGrooveRight)
    print("pointOnGrooveRight：", pointOnGrooveRight.shape)

    # ***************2.对右侧坡口上的图片点再进行Hough直线检测，返回该直线参数******************
    gray = cv.cvtColor(pointOnGrooveRight, cv.COLOR_BGR2GRAY)  # 先转换为灰度图
    _, binary = cv.threshold(gray, 100, 255, cv.THRESH_BINARY)  # 必须是单通道灰度图像

    linesOnGrooveRight = cv.HoughLines(binary, 1, np.pi / 180, 14)  # 返回一个：math:(rho，theta)值的数组
    print(linesOnGrooveRight.shape)  # (1,1,2)
    print("linesOnGrooveLeft(rho，theta):", linesOnGrooveRight)  # 获取主线参数1

    for lineRight in linesOnGrooveRight:
        rho, theta = lineRight[0]
        print("右半部坡口主线角度：", theta * 180 / 3.14)

        a, b = np.cos(theta), np.sin(theta)
        x0, y0 = a * rho, b * rho  # 原点到直线的 垂点坐标
        lenL, lenR = 500, 600
        x1, y1 = int(x0 + lenL * (-b)), int(y0 + lenL * (a))
        x2, y2 = int(x0 - lenR * (-b)), int(y0 - lenR * (a))
        # 获取主线参数2
        k = (y2 - y1) / (x2 - x1)
        b = y1 - k * x1
        print("右边半部坡口斜率k: %f" % k, "截距b：%f" % b)
        cv.line(img_new_Right, (x1, y1), (x2, y2), (0, 255, 0), 1)
    cv.imshow("mainline of pointOnGrooveRight", pointOnGrooveRight)  # 显示不了主线，则取到主线参数也可以
    return linesOnGrooveRight, k, b


def featurePoint(img, k1, b1, k2, b2):
    x = (b2 - b1) / (k1 - k2)  # 列坐标
    y = k1 * x + b1  # 行坐标
    # a, b = int(x), int(y)  # 直接省略小数部分
    row, col = round(y), round(x)
    img[row, col] = [0, 255, 0]

    cv.imshow("FeaturePoint", img)
    cv.imwrite(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\FeaturePoint_Src.png", img)
    return row, col


src = cv.imread(r"D:\SHU\Research Group\LaserVisionSensor\LaserImage\MVS_Data\20220909\V\Image_20220909101553637.jpg")
src = cv.resize(src, None, fx=0.4, fy=0.4, interpolation=cv.INTER_CUBIC)

img_new, CP_row, CP_col = gray_center_method(src)

img_new_Left, img_new_Right, furthestPoint_row_Left, furthestPoint_col_Left, furthestPoint_row_Right, furthestPoint_col_Right = HoughLines(
    img_new, CP_row, CP_col)
linesOnGrooveRight, k1, b1 = HoughLines_Left(img_new_Left, furthestPoint_row_Left, furthestPoint_col_Left)
linesOnGrooveRight, k2, b2 = HoughLines_Right(img_new_Right, furthestPoint_row_Right, furthestPoint_col_Right)
featurePoint(src, k1, b1, k2, b2)

cv.waitKey(0)
cv.destroyAllWindows()
